import cv2
import numpy as np
import math,random
from sklearn.cluster import DBSCAN
from scipy.stats import binned_statistic
from sklearn.linear_model import RANSACRegressor
from sklearn.linear_model import LinearRegression
from sklearn.cluster import MeanShift  # 比DBSCAN更快

class LineClassifier:
    """
    线段分类器，策略模式实现不同分类算法
    """
    def classify(self, lines, outlier_threshold=20):
        """
        分类线段为横向和纵向
        """
        angles = self._calculate_angles(lines)
        median_angle = np.median(angles)
        lower_bound = median_angle - 45
        upper_bound = median_angle + 45
        
        # 处理角度边界
        if lower_bound < -90:
            upper_bound += 90
            lower_bound += 90
        elif upper_bound > 90:
            lower_bound -= 90
            upper_bound -= 90
            
        # 初步分类
        horizontal, vertical = [], []
        for line, angle in zip(lines, angles):
            if lower_bound <= angle <= upper_bound:
                horizontal.append({'points': line[0], 'angle': angle})
            else:
                vertical.append({'points': line[0], 'angle': angle})
        
        # 过滤离群值
        h_inliers, h_outliers, h_mode_angle = self._filter_outliers(horizontal, True)
        v_mode_angle = h_mode_angle - 90
        if v_mode_angle < -90:
            v_mode_angle += 180
        v_inliers, v_outliers, _ = self._filter_outliers(vertical, False, v_mode_angle)
        
        # 计算平均角度
        avg_h_angle = np.median([line['angle'] for line in h_inliers])
        avg_v_angle = np.median([line['angle'] for line in v_inliers])
        
        # 二次分类离群值
        for line in h_outliers + v_outliers:
            angle = line['angle']
            dist_h = self._angle_diff(angle, avg_h_angle)
            dist_v = self._angle_diff(angle, avg_v_angle)
            
            if dist_h - outlier_threshold > 0 and dist_v - outlier_threshold > 0:
                continue
                
            if dist_h < dist_v:
                h_inliers.append(line)
            else:
                v_inliers.append(line)
        
        horizontal_lines = [line['points'] for line in h_inliers]
        vertical_lines = [line['points'] for line in v_inliers]
        
        return horizontal_lines, vertical_lines, avg_h_angle, avg_v_angle
    
    def _calculate_angles(self, lines):
        """
        计算线段角度
        """
        angles = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            dx, dy = x2 - x1, y2 - y1
            angle = math.degrees(math.atan2(dy, dx))
            # 归一化角度到 [-90, 90]
            if angle < -90:
                angle += 180
            elif angle > 90:
                angle -= 180
            angles.append(angle)
        return angles
    
    def _filter_outliers(self, line_list, flag=True, reference_angle=None, bin_size=10, outlier_threshold=20):
        """
        过滤离群线段
        """
        if not line_list:
            return [], [], None
        if flag:
            angles = np.array([line['angle'] for line in line_list])
            # 计算区间众数
            bin_edges = np.arange(-90, 90 + bin_size, bin_size)
            bin_counts, _, _ = binned_statistic(angles, angles, statistic='count', bins=bin_edges)
            max_bin_idx = np.argmax(bin_counts)
            mode_lower = bin_edges[max_bin_idx]
            mode_upper = bin_edges[max_bin_idx + 1]
            mode_angle = (mode_lower + mode_upper) / 2
        else:
            mode_angle = reference_angle
            
        # 筛选离群值
        inliers = [line for line in line_list if abs(line['angle'] - mode_angle) <= outlier_threshold]
        outliers = [line for line in line_list if abs(line['angle'] - mode_angle) > outlier_threshold]
        
        return inliers, outliers, mode_angle
    
    def _angle_diff(self, a, b):
        """
        计算角度差值
        """
        diff = abs(a - b)
        return min(diff, 180 - diff)


class LineFitter:
    def __init__(self):
        self.ransac = RANSACRegressor(
                estimator=LinearRegression(),
                min_samples=3,
                residual_threshold=12.0,
                random_state=0,
            )

    def fit(self, lines, avg_angle, depth_img):
        """二次拟合+离群点剔除的完整实现"""
        fitted = []
        if not lines:
            return fitted
        
        # vis_img = depth_img.copy() if depth_img.ndim == 3 else cv2.cvtColor(depth_img, cv2.COLOR_GRAY2BGR)
        # 数据准备
        lines_arr = np.array([[x1, y1, x2, y2] for x1, y1, x2, y2 in lines], dtype=np.float32)
        mid_points = (lines_arr[:, :2] + lines_arr[:, 2:]) / 2
        # 清理无效数据（NaN/Inf）
        valid_mask = np.all(np.isfinite(mid_points), axis=1)
        mid_points = mid_points[valid_mask]
        lines_arr = lines_arr[valid_mask]
        # 检查数据是否为空
        if len(mid_points) == 0:
            return fitted  # 直接返回空结果
        # 确保角度有效
        if not np.isfinite(avg_angle):
            avg_angle = 0.0  # 默认值
        # 计算投影轴（增加数值稳定性检查）
        rad_h = np.deg2rad(avg_angle)
        proj_axis = np.array([-np.sin(rad_h), np.cos(rad_h)])
        if not np.all(np.isfinite(proj_axis)):
            proj_axis = np.array([0.0, 1.0])  # 默认垂直方向
        # 计算投影值
        projections = mid_points.dot(proj_axis).reshape(-1, 1)
   
        # 聚类
        dbscan = DBSCAN(eps=20, min_samples=3)
        clusters = dbscan.fit_predict(projections)
        unique_clusters = np.unique(clusters)

        for cluster_id in unique_clusters:
            if cluster_id == -1:
                continue
                
            # 获取当前聚类的所有点
            cluster_lines = lines_arr[clusters == cluster_id]
            all_points = np.vstack([cluster_lines[:,:2], cluster_lines[:,2:]])
            points = np.unique(all_points, axis=0)

            if len(points) < 5:  # 至少需要5个点进行二次拟合
                continue        
            # 第一次拟合
            X = points[:, 0].reshape(-1, 1)
            y = points[:, 1]
            try:
                # ========== 第一次拟合 ==========
                self.ransac.fit(X, y)
                if not hasattr(self.ransac, "estimator_"):
                    print("RANSAC failed to fit a model. Skipping...")
                    continue
                # 获取结果
                inlier_mask = self.ransac.inlier_mask_
                k1 = self.ransac.estimator_.coef_[0]
                b1 = self.ransac.estimator_.intercept_
                
                # ========== 剔点逻辑 ==========
                distances = np.abs(points[:,1] - (k1*points[:,0] + b1)) / np.sqrt(k1**2 + 1)
                # 动态阈值：使用IQR方法自动确定
                q1, q3 = np.percentile(distances, [25, 75])
                iqr = q3 - q1
                threshold = min(20,iqr)  # 硬上限20像素
                # 保留内点
                inlier_mask = distances <= threshold
                clean_points = points[inlier_mask]
                
                # # ========== 第二次拟合 ==========
                if len(clean_points) > 4:  # 至少保留2个点
                    X_clean = clean_points[:, 0].reshape(-1, 1)
                    y_clean = clean_points[:, 1]
                    lr2 = LinearRegression().fit(X_clean, y_clean) 
                    k, b = lr2.coef_[0], lr2.intercept_
                    fitted.append((k, b))
                            
            except Exception as e:
                print(f"Line fitting error: {e}")
                continue
                
        return fitted



class IntersectionFinder:
    """
    交点查找器
    """
    def find(self, h_lines, v_lines, w, h, margin=30):
        """
        查找交点
        """
        intersections = []
        for h_line in h_lines:
            for v_line in v_lines:
                point = self._find_intersection(h_line, v_line, w, h, margin)
                if point is not None:
                    intersections.append(point)
        return intersections
    
    def _find_intersection(self, line1, line2, w, h, margin=30):
        """
        计算两条直线的交点
        """
        k1, b1 = line1
        k2, b2 = line2
        
        if abs(k1 - k2) < 1e-6:
            return None
            
        x = (b2 - b1) / (k1 - k2)
        y = k1 * x + b1
        if (margin <= x < w - margin) and (margin <= y < h - margin):
            return (int(x), int(y))
        else:
            return None


class MaskCreator:
    """
    掩膜创建器
    """
    def create(self, image, points, radius=30):
        """
        创建掩膜图像
        """
        mask = np.zeros_like(image)
        for (x, y) in points:
            cv2.circle(mask, (x, y), radius, 255, -1) 
        masked_img = cv2.bitwise_and(image, mask)
        return masked_img


class ResultVisualizer:
    """
    结果可视化器
    """
    def draw_lines(self, img, h_lines, v_lines):
        """
        绘制拟合的直线
        """
        vis_img = img.copy()
        h, w = img.shape[:2]
        
        # 绘制横向直线
        for k, b in h_lines:
            self._draw_line(vis_img, k, b, (255, 255, 255))
            
        # 绘制纵向直线
        for k, b in v_lines:
            self._draw_line(vis_img, k, b, (255, 255, 255))
            
        return vis_img
    
    def draw_intersections(self, img, intersections):
        """
        绘制交点
        """
        vis_img = img.copy()
        for (x, y) in intersections:
            cv2.circle(vis_img, (x, y), 5, (0, 0, 255), 5)
        return vis_img
    
    def _draw_line(self, img, k, b, color):
        """
        绘制单条直线
        """
        h, w = img.shape[:2]
        
        if abs(k) < 1e4:
            y1 = k*0 + b
            y2 = k*(w-1) + b
            
            if y1 < 0 or y1 >= h or y2 < 0 or y2 >= h:
                points = []
                if k != 0:
                    x_top = (0 - b) / k
                    if 0 <= x_top <= w-1:
                        points.append((x_top, 0))
                x_bottom = (h-1 - b) / k
                if 0 <= x_bottom <= w-1:
                    points.append((x_bottom, h-1))
                y_left = b
                if 0 <= y_left <= h-1:
                    points.append((0, y_left))
                y_right = k*(w-1) + b
                if 0 <= y_right <= h-1:
                    points.append((w-1, y_right))
                
                if len(points) >= 2:
                    if len(points) > 2:
                        max_dist = 0
                        best_pair = None
                        for i in range(len(points)):
                            for j in range(i+1, len(points)):
                                dist = np.linalg.norm(np.array(points[i]) - np.array(points[j]))
                                if dist > max_dist:
                                    max_dist = dist
                                    best_pair = (points[i], points[j])
                        if best_pair:
                            pt1, pt2 = best_pair
                        else:
                            pt1, pt2 = points[0], points[1]
                    else:
                        pt1, pt2 = points[0], points[1]
                    
                    cv2.line(img, 
                            (int(pt1[0]), int(pt1[1])), 
                            (int(pt2[0]), int(pt2[1])), 
                            color, 3)
            else:
                cv2.line(img, (0, int(y1)), (w-1, int(y2)), color, 3)
        else:
            x_val = -b / k
            cv2.line(img, (int(x_val), 0), (int(x_val), h-1), color, 3)

    def _display_results(self, result):
        """
        显示结果
        """
        if result is None:
            return
            
        cv2.imshow("Original", result['original'])
        cv2.imshow("Edges", result['edges'])
        
        # 绘制原始线段
        l_img = np.zeros((result['original'].shape[0], result['original'].shape[1], 3), 
                         dtype=result['original'].dtype)
        for line in result['lines']:
            x1, y1, x2, y2 = map(int, line[0])
            cv2.line(l_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.imshow("Detected Lines", l_img)
        
        cv2.imshow("Fitted Rebars", result['fitted'])
        cv2.imshow("Masked Image", result['masked'])
        
        c = cv2.waitKey(0)
        if c == 113:  # 'q'键
            cv2.destroyAllWindows()
            exit(0)
            return False
        return True